import cv2
import time
import numpy as np

from cvs import cvs

from yolov5 import Yolov5
from deep_sort import DeepSort
from utils import draw_track_res
from mytracker import MyTracker
from utils import detect_postprocess, preprocess_img, draw_detect_res

from qiniu_test import upload_qiniuyun
from miaotixing_test import miaotixing_post


video_path = "/home/aidlux_deepsort_simpletracker/TownCentreXVID.avi"
track_model_path = "/home/aidlux_deepsort_simpletracker/models/feature_extractor_osnet.tflite"

# track_model_path = "./models/feature_extractor_osnet_int8.tflite"

# detector = Yolov5(model_path="/home/aidlux_deepsort_simpletracker/models/yolov5n-640.tflite", 
#                     input_shape=(1, 640, 640, 3),
#                     output_shape=(1, 25200, 85),
#                     device=0,
#                     int8=False,
#                     input_scale=None,
#                     input_zero_point=None,
#                     output_sacle=None,
#                     output_zero_point=None,
#                     conf_thres=0.4, 
#                     iou_thres=0.4, 
#                     interested_classid=None)

# tracktor = DeepSort(model_path="/home/aidlux_deepsort_simpletracker/models/feature_extractor_osnet.tflite", 
#                     input_shape=(1, 64, 64, 3),
#                     output_shape=(1, 512),
#                     device=0,
#                     int8=False,
#                     input_scale=None,
#                     input_zero_point=None,
#                     output_sacle=None,
#                     output_zero_point=None,
#                     max_dist=0.2, 
#                     max_iou_distance=0.7, 
#                     max_age=30, n_init=3, nn_budget=100)



detector = Yolov5(model_path="/home/aidlux_deepsort_simpletracker/models/yolov5n-640.tflite", 
                    input_shape=(1, 640, 640, 3),
                    output_shape=(1, 25200, 85),
                    device=0,
                    int8=False,
                    input_scale=None,
                    input_zero_point=None,
                    output_sacle=None,
                    output_zero_point=None,
                    conf_thres=0.4, 
                    iou_thres=0.4, 
                    interested_classid=None)
                    
tracktor = MyTracker(model_path="/home/aidlux_deepsort_simpletracker/models/feature_extractor_osnet.tflite", 
                        input_shape=(1, 64, 64, 3),
                        output_shape=(1, 512),
                        device=0,
                        int8=False,
                        input_scale=None,
                        input_zero_point=None,
                        output_sacle=None,
                        output_zero_point=None,
                        max_age=40,
                        momentum=0.9)


# tracktor = MyTracker(model_path="/home/aidlux_deepsort_simpletracker/models/feature_extractor_osnet_int8.tflite", 
#                         input_shape=(1, 64, 64, 3),
#                         output_shape=(1, 512),
#                         device=0,
#                         int8=True,
#                         input_scale=0.01865844801068306,
#                         input_zero_point=114,
#                         output_sacle=0.07101853936910629,
#                         output_zero_point=None,
#                         max_age=40,
#                         momentum=0.9)

# detector = Yolov5(model_path="/home/aidlux_deepsort_simpletracker/models/yolov5n_320_int8.tflite", 
#                     input_shape=(1, 320, 320, 3),
#                     output_shape=(1, 6300, 85),
#                     device=0,
#                     int8=True,
#                     input_scale=None,
#                     input_zero_point=None,
#                     output_sacle=0.004703259561210871,
#                     output_zero_point=3,
#                     conf_thres=0.4, 
#                     iou_thres=0.4, 
#                     interested_classid=[0])

if __name__ == '__main__':
    print("start............ \n")

    # 根据视频尺寸，填充一个polygon，供撞线计算使用
    mask_image_temp = np.zeros((1080, 1920), dtype=np.uint8)

    # 初始化2个撞线polygon
    list_pts_blue = [[204, 305], [227, 431], [605, 522], [1101, 464], [1900, 601], [1902, 495], [1125, 379], [604, 437],
                     [299, 375], [267, 289]]
    ndarray_pts_blue = np.array(list_pts_blue, np.int32)
    polygon_blue_value_1 = cv2.fillPoly(mask_image_temp, [ndarray_pts_blue], color=1)
    polygon_blue_value_1 = polygon_blue_value_1[:, :, np.newaxis]

    # 填充第二个polygon
    mask_image_temp = np.zeros((1080, 1920), dtype=np.uint8)
    list_pts_yellow = [[181, 305], [207, 442], [603, 544], [1107, 485], [1898, 625], [1893, 701], [1101, 568],
                       [594, 637], [118, 483], [109, 303]]
    ndarray_pts_yellow = np.array(list_pts_yellow, np.int32)
    polygon_yellow_value_2 = cv2.fillPoly(mask_image_temp, [ndarray_pts_yellow], color=2)
    polygon_yellow_value_2 = polygon_yellow_value_2[:, :, np.newaxis]

    # 撞线检测用mask，包含2个polygon，（值范围 0、1、2），供撞线计算使用
    polygon_mask_blue_and_yellow = polygon_blue_value_1 + polygon_yellow_value_2

    # 缩小尺寸，1920x1080->960x540
    polygon_mask_blue_and_yellow = cv2.resize(polygon_mask_blue_and_yellow, (720, 480))

    # 蓝 色盘 b,g,r
    blue_color_plate = [255, 0, 0]
    # 蓝 polygon图片
    blue_image = np.array(polygon_blue_value_1 * blue_color_plate, np.uint8)

    # 黄 色盘
    yellow_color_plate = [0, 255, 255]
    # 黄 polygon图片
    yellow_image = np.array(polygon_yellow_value_2 * yellow_color_plate, np.uint8)

    # 彩色图片（值范围 0-255）
    color_polygons_image = blue_image + yellow_image
    # 缩小尺寸，1920x1080->960x540
    color_polygons_image = cv2.resize(color_polygons_image, (720, 480))

    # list 与蓝色polygon重叠
    list_overlapping_blue_polygon = []

    # list 与黄色polygon重叠
    list_overlapping_yellow_polygon = []

    # 进入数量
    down_count = 0
    # 离开数量
    up_count = 0

    down_flag = False
    up_flag = False
    

    font_draw_number = cv2.FONT_HERSHEY_SIMPLEX
    draw_text_postion = (int(720 * 0.01), int(480 * 0.05))



    # 打开视频
    capture = cv2.VideoCapture("/home/aidlux_deepsort_simpletracker/TownCentreXVID.avi")


    fps_num = 0
    save_abnormal_image = True

    while True:
        # 读取每帧图片
        # _, im = capture.read()
        # if im is None:
        #     break


        # frame = capture.read()
        ret, frame = capture.read()
        fps_num += 1
        if fps_num % 15 == 0:            
            fps_num = 0
        else:
            continue
            
        if frame is None:
            # print(f"video_path:{video_path} is done~")
            continue
        
        frame = cv2.resize(frame, (720, 480))

        det_time = time.time()

        

        # det_res = detector.invoke(frame)
        det_res = detector.invoke(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))

        det_time = time.time() - det_time
        if len(det_res) == 0:
            continue

        x1y1x2y2, confs, clss = det_res[:, 0:4], det_res[:, 4], det_res[:, 5]
        
        track_time = time.time()
        # outputs = tracktor.update(x1y1x2y2, confs, clss, frame)
        outputs = tracktor(x1y1x2y2, confs, clss, frame)
        track_time = time.time() - track_time

        if len(outputs) > 0:
            frame = draw_track_res(frame, outputs)
        else:
            frame = detector.draw_pic(frame, det_res)
        # cv2.putText(frame, "Track Time:{:^2.0f}ms, FPS:{:^2.0f}".format(track_time*1000, 1/track_time), (2, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 1)
        # cv2.putText(frame, "Detect Time:{:^2.0f}ms, FPS:{:^2.0f}".format(det_time*1000, 1/det_time), (2, 44), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 1)
        # cvs.imshow(frame)
    

        # 输出图片
        output_image_frame = cv2.add(frame, color_polygons_image)

        if len(outputs) > 0:
            # ----------------------判断撞线----------------------
            for item_bbox in outputs:
                x1, y1, x2, y2, track_id, label,score = item_bbox

                # 撞线检测点，(x1，y1)，y方向偏移比例 0.0~1.0
                y1_offset = int(y1 + ((y2 - y1) * 0.6))

                # 撞线的点
                y = int(y1_offset)
                x = int(x1)

                if polygon_mask_blue_and_yellow[y, x] == 1:
                    # 如果撞 蓝polygon
                    if track_id not in list_overlapping_blue_polygon:
                        list_overlapping_blue_polygon.append(track_id)
                    pass

                    # 判断 黄polygon list 里是否有此 track_id
                    # 有此 track_id，则 认为是 外出方向
                    if track_id in list_overlapping_yellow_polygon:
                        # 外出+1
                        up_count += 1
                        up_flag = True

                        print(f'类别: {label} | id: {track_id} | 上行撞线 | 上行撞线总数: {up_count} | 上行id列表: {list_overlapping_yellow_polygon}')


                        # 删除 黄polygon list 中的此id
                        list_overlapping_yellow_polygon.remove(track_id)

                        pass
                    else:
                        # 无此 track_id，不做其他操作
                        pass

                elif polygon_mask_blue_and_yellow[y, x] == 2:
                    # 如果撞 黄polygon
                    if track_id not in list_overlapping_yellow_polygon:
                        list_overlapping_yellow_polygon.append(track_id)
                    pass

                    # 判断 蓝polygon list 里是否有此 track_id
                    # 有此 track_id，则 认为是 进入方向
                    if track_id in list_overlapping_blue_polygon:
                        # 进入+1
                        down_count += 1
                        down_flag = True     

                        print(f'类别: {label} | id: {track_id} | 下行撞线 | 下行撞线总数: {down_count} | 下行id列表: {list_overlapping_blue_polygon}')

                        # 删除 蓝polygon list 中的此id
                        list_overlapping_blue_polygon.remove(track_id)

                        pass
                    else:
                        # 无此 track_id，不做其他操作
                        pass
                    pass
                else:
                    pass
                pass

            pass

            # ----------------------清除无用id----------------------
            list_overlapping_all = list_overlapping_yellow_polygon + list_overlapping_blue_polygon
            for id1 in list_overlapping_all:
                is_found = False
                for _, _, _, _, bbox_id, _,_ in outputs:
                    if bbox_id == id1:
                        is_found = True
                        break
                    pass
                pass

                if not is_found:
                    # 如果没找到，删除id
                    if id1 in list_overlapping_yellow_polygon:
                        list_overlapping_yellow_polygon.remove(id1)
                    pass
                    if id1 in list_overlapping_blue_polygon:
                        list_overlapping_blue_polygon.remove(id1)
                    pass
                pass
            list_overlapping_all.clear()
            pass

            # 清空list
            # outputs.clear()

            pass
        else:
            # 如果图像中没有任何的bbox，则清空list
            list_overlapping_blue_polygon.clear()
            list_overlapping_yellow_polygon.clear()
            pass
        pass

        text_draw = 'DOWN: ' + str(down_count) + \
                    ' , UP: ' + str(up_count)
        output_image_frame = cv2.putText(img=output_image_frame, text=text_draw,
                                         org=draw_text_postion,
                                         fontFace=font_draw_number,
                                         fontScale=1, color=(255, 255, 255), thickness=2)

        # cv2.imshow('demo', output_image_frame)        
        # cv2.waitKey(1)


        if save_abnormal_image:

            if up_flag:
                up_flag = False                
                img_save_path = 'up_count_abnormal_{}.jpg'.format(up_count)                            
                cv2.imwrite(img_save_path, output_image_frame)
                img_name = img_save_path.split('\\')[-1]
                url_receive = upload_qiniuyun(img_name, img_save_path)
                from miaotixing_test import miaotixing_post
                miaotixing_result = miaotixing_post(url_receive)

            if down_flag:
                down_flag = False                
                img_save_path = 'down_count_abnormal_{}.jpg'.format(down_count)                            
                cv2.imwrite(img_save_path, output_image_frame)
                img_name = img_save_path.split('\\')[-1]
                url_receive = upload_qiniuyun(img_name, img_save_path)
                from miaotixing_test import miaotixing_post
                miaotixing_result = miaotixing_post(url_receive)


        cvs.imshow(output_image_frame)

        pass
    pass

    # capture.release()
    # cv2.destroyAllWindows()
